import codecs
import os
import cv2 as cv
from zipfile import ZipFile
from skimage import io
import numpy as np
import random
import logging
import time
import yaml


def random_transform(img, scale_min=0.5, scale_max=1.5):
    h, w = img.shape[:2]

    scale = np.random.uniform(scale_min, scale_max)
    scaled_img = cv.resize(img, None, fx=scale, fy=scale)
    h_s, w_s = scaled_img.shape[:2]
    h_start = max(0, int((h_s - h) / 2))
    w_start = max(0, int((w_s - w) / 2))
    cropped_img = scaled_img[h_start:h_start + h, w_start:w_start + w]

    angle = np.random.uniform(-45, 45)
    M = cv.getRotationMatrix2D((w / 2, h / 2), angle, 1)
    rotated_img = cv.warpAffine(cropped_img, M, (w, h))

    return rotated_img


def process_images_from_zip(zip_path, scale_min=0.5, scale_max=1.5, num_images=10, progress_callback=None,
                            thread_event=None, output_type='RGB'):
    out_path = 'out/'
    if not os.path.exists(out_path):
        logging.info('创建输出目录 %s', out_path)
        os.mkdir(out_path)

    logging.info('开始处理zip文件中的图像')
    with ZipFile(zip_path, 'r') as zipped:
        namelist = [name for name in zipped.namelist() if
                    name.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif'))]
        length = len(namelist)

        for name in namelist:
            if progress_callback:
                progress_callback(int(100 * (namelist.index(name) + 1) / length))

            if thread_event and thread_event.is_set():
                logging.info('用户点击了取消按钮')
                return

            with zipped.open(name) as file:
                if output_type == 'RGB':
                    img = cv.imdecode(np.frombuffer(file.read(), np.uint8), 1)
                elif output_type == 'Grayscale':
                    img = cv.imdecode(np.frombuffer(file.read(), np.uint8), 0)
                elif output_type == 'Binary':
                    img = cv.imdecode(np.frombuffer(file.read(), np.uint8), 0)
                    thresh_hold, img = cv.threshold(img, 0, 255, cv.THRESH_OTSU)
                else:
                    raise ValueError('无效的输出类型: ' + output_type)

                new_img = cv.resize(img, (224, 224), interpolation=cv.INTER_NEAREST)

                output_name = os.path.join(out_path, os.path.basename(name))
                logging.info(f'保存处理后的图像到: {output_name}')

                io.imsave(output_name, new_img)

                for i in range(num_images - 1):
                    transformed_img = random_transform(new_img)
                    output_name_transformed = os.path.join(out_path,
                                                           f"{os.path.splitext(os.path.basename(name))[0]}_transformed_{i}.jpg")
                    io.imsave(output_name_transformed, transformed_img)


def load_config():
    with open('config.yaml', 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    return config